import csv


class TrainTemplate:
    ''' 列车模板 '''

    Code = ""
    EventList = None

    def __init__(self) -> None:
        self.EventList = list()


class TrainTemplateEvent:
    ''' 列车模板事件 '''
    Seq = -1
    MovementType = ""
    RouteTag = ""
    MinDuration = -1
    MaxDuration = -1


class Train:
    ''' 列车 '''
    ID = ""
    Template = None
    Seq = -1


TrainTemplateDict = dict()
DataDict = dict()

# 载入列车模板
f = open("./LocalData/TestBed_TrainTemplate.csv", "r")
for item in csv.DictReader(f):
    templateCode = item["TemplateCode"]
    if templateCode in TrainTemplateDict:
        trainTemplate = TrainTemplateDict[templateCode]
    else:
        trainTemplate = TrainTemplate()
        trainTemplate.Code = templateCode
        TrainTemplateDict[templateCode] = trainTemplate
    event = TrainTemplateEvent()
    event.Seq = int(item["Seq"])
    event.MovementType = item["MovementType"]
    event.RouteTag = item["RouteTag"]
    event.MinDuration = item["MinDuration"]
    event.MaxDuration = item["MaxDuration"]
    trainTemplate.EventList.append(event)
f.close()


# 载入列车比例方案
f = open("./LocalData/TestBed_TrainCombScheme.csv", "r")
for item in csv.DictReader(f):
    description = item["Description"]
    trainDict = dict()

    for templateCode in TrainTemplateDict:
        num = int(item[templateCode])
        template = TrainTemplateDict[templateCode]
        for i in range(0, num):
            train = Train()
            train.Template = template
            train.ID = templateCode + str(i+1)
            train.Seq = i+1
            trainDict[train.ID] = train

    DataDict[description] = trainDict
f.close()

print("Finished!")


# 生成列车比例方案数据集
for desc in DataDict:
    fTrain = open(
        "./LocalData/TestBed/Train_{0}.csv".format(desc), "w", encoding="utf-8", newline="")
    fieldNames = ["ID", "Type", "Seq"]
    trainWriter = csv.DictWriter(fTrain, fieldNames)
    trainWriter.writeheader()
    for trID in DataDict[desc]:
        train = DataDict[desc][trID]
        trainWriter.writerow({"ID": train.ID, "Type": "G", "Seq":train.Seq})
    f.close()

    fTimetable = open(
        "./LocalData/TestBed/Timetable_{0}.csv".format(desc), "w", encoding="utf-8", newline="")
    fieldNames = ["StationID", "ID", "TrainID", "MovementType", "Sequence", "RouteTag",
                  "StartTime", "EndTime", "MinDuration", "MaxDuration"]
    timetableWriter = csv.DictWriter(fTimetable, fieldNames)
    timetableWriter.writeheader()
    idx = 0
    for trID in DataDict[desc]:
        train = DataDict[desc][trID]
        for event in train.Template.EventList:
            timetableWriter.writerow({"StationID": "ZDD", "ID": idx, "TrainID": train.ID, "MovementType": event.MovementType, "Sequence": event.Seq, "RouteTag": event.RouteTag, "StartTime": "00:00:00", "EndTime": "00:00:00", "MinDuration": event.MinDuration, "MaxDuration": event.MaxDuration})
            idx += 1
    fTimetable.close()
